/*
 * Copyright (C) 2020 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */

package vip.justlive.easyproxy.core.httpproxy;

import java.net.HttpURLConnection;
import java.net.URI;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Objects;
import lombok.extern.slf4j.Slf4j;
import vip.justlive.easyproxy.core.context.Config;
import vip.justlive.easyproxy.core.context.Metrics;
import vip.justlive.oxygen.core.config.ConfigFactory;
import vip.justlive.oxygen.core.exception.Exceptions;
import vip.justlive.oxygen.core.net.aio.core.AioListener;
import vip.justlive.oxygen.core.net.aio.core.ChannelContext;
import vip.justlive.oxygen.core.net.aio.core.Client;
import vip.justlive.oxygen.core.net.aio.core.GroupContext;
import vip.justlive.oxygen.core.net.http.HttpMethod;
import vip.justlive.oxygen.core.util.Bytes;
import vip.justlive.oxygen.core.util.HttpHeaders;
import vip.justlive.oxygen.core.util.Strings;
import vip.justlive.oxygen.web.http.Request;
import vip.justlive.oxygen.web.http.Response;
import vip.justlive.oxygen.web.server.aio.HttpServerAioHandler;

/**
 * http代理服务端处理
 *
 * @author wubo
 */
@Slf4j
public class HttpServerHandler extends HttpServerAioHandler implements AioListener {

  private static final String HTTPS = "https";
  private static final byte[] HTTPS_RES = "HTTP/1.1 200 Connection Established\r\n\r\n".getBytes();

  private final Config config;
  private final Client client;

  public HttpServerHandler(Config config) {
    super(Strings.EMPTY);
    this.config = config;
    HttpClientHandler handler = new HttpClientHandler();
    GroupContext context = new GroupContext(handler).setAioListener(handler);
    ConfigFactory.load(context, Config.HTTP_PROXY_AIO_CONFIG);
    client = new Client(context);
  }

  @Override
  public Object decode(ByteBuffer buffer, int readableSize, ChannelContext channelContext) {
    int index = buffer.position();
    if (channelContext.getAttr(HTTPS) != null) {
      byte[] bytes = new byte[readableSize];
      buffer.get(bytes);
      handleHttps(bytes, channelContext);
      return bytes;
    }
    Request req = (Request) super.decode(buffer, readableSize, channelContext);
    if (req != null) {
      Bytes bytes = new Bytes();
      bytes.write(buffer.array(), index, buffer.position());
      handleHttp(req, bytes, channelContext);
    }
    return req;
  }

  @Override
  public void handle(Object data, ChannelContext channelContext) {
  }

  @Override
  public ByteBuffer encode(Object data, ChannelContext channelContext) {
    byte[] bytes = (byte[]) data;
    Metrics.metrics(channelContext.getServerAddress().getPort()).getWriteBytes()
        .getAndAdd(bytes.length);
    return ByteBuffer.wrap(bytes);
  }

  @Override
  public void onConnected(ChannelContext channelContext) {
    channelContext.addAttr(Config.REMOTE, new ArrayList<>());
    Metrics.metrics(channelContext.getServerAddress().getPort()).getServerChannels()
        .set(channelContext.getGroupContext().getChannels().size());
  }

  @Override
  public void onClosed(ChannelContext channelContext) {
    Metrics.metrics(channelContext.getServerAddress().getPort()).getServerChannels()
        .set(channelContext.getGroupContext().getChannels().size());
    List<ChannelContext> channels = clientChannels(channelContext);
    channels.forEach(ChannelContext::close);
    channels.clear();
    Metrics.metrics(channelContext.getServerAddress().getPort()).getClientChannels()
        .set(client.getGroupContext().getChannels().size());
  }

  @SuppressWarnings("unchecked")
  private List<ChannelContext> clientChannels(ChannelContext channelContext) {
    return (List<ChannelContext>) channelContext.getAttr(Config.REMOTE);
  }

  private void handleHttps(byte[] bytes, ChannelContext channelContext) {
    Metrics.metrics(channelContext.getServerAddress().getPort()).getReadBytes()
        .getAndAdd(bytes.length);
    List<ChannelContext> channels = clientChannels(channelContext);
    if (channels.size() != 1) {
      log.warn("https请求连接客户端连接数量有误, {}", channels);
      channelContext.close();
      return;
    }
    channels.get(0).write(bytes);
  }

  private void handleHttp(Request request, Bytes bytes, ChannelContext channelContext) {
    Metrics.metrics(channelContext.getServerAddress().getPort()).getReadBytes()
        .getAndAdd(bytes.size());
    if (config.getHttpProxy().isAuth() && !basicAuth(request)) {
      channelContext.write(badAuth(request));
      return;
    }
    String url = request.getRequestUri();
    if (request.getMethod() == HttpMethod.CONNECT) {
      url = HttpHeaders.HTTPS_PREFIX + url;
    }
    try {
      URI uri = new URI(url);
      ChannelContext channel = client
          .connect(uri.getHost(), uri.getPort() == -1 ? 80 : uri.getPort());
      channel.addAttr(Config.REMOTE, channelContext);
      channel.join();
      clientChannels(channelContext).add(channel);
      if (request.getMethod() == HttpMethod.CONNECT) {
        channelContext.addAttr(HTTPS, HTTPS);
        channelContext.write(HTTPS_RES);
      } else {
        channel.write(bytes.toArray());
      }
      Metrics.metrics(channelContext.getServerAddress().getPort()).getClientChannels()
          .set(client.getGroupContext().getChannels().size());
    } catch (Exception e) {
      log.error("connect uri {} failed", url);
      throw Exceptions.wrap(e);
    }
  }

  private boolean basicAuth(Request request) {
    String auth = request.getHeader("Proxy-Authorization");
    if (!Strings.hasText(auth)) {
      return false;
    }
    auth = new String(Base64.getDecoder().decode(auth.substring(6)));
    String[] up = auth.split(Strings.COLON, 2);
    return Objects.equals(ConfigFactory.getProperty(Config.HTTP_PROXY_AUTH_PREFIX + up[0]), up[1]);
  }

  private byte[] badAuth(Request request) {
    Response response = new Response(request);
    response.setStatus(HttpURLConnection.HTTP_PROXY_AUTH);
    response.setHeader("Proxy-Authenticate", "Basic realm=\"Text\"");
    return super.encode(response, null).array();
  }
}
